"""
高级运动控制算法
实现六足机器人的高级运动控制功能
"""

import math
import numpy as np
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
from .hexapod_controller import HexapodController


@dataclass
class LegState:
    """腿部状态"""

    position: Tuple[float, float, float]  # 末端位置 (x, y, z)
    velocity: Tuple[float, float, float]  # 末端速度
    force: Tuple[float, float, float]  # 末端受力
    contact: bool  # 是否接触地面


class AdvancedMotionController(HexapodController):
    """高级运动控制器"""

    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)

        # 运动参数
        self.step_height = 50  # 抬腿高度
        self.step_length = 100  # 步长
        self.step_period = 1.0  # 步态周期
        self.body_height = 150  # 身体高度
        self.body_roll = 0  # 身体横滚角
        self.body_pitch = 0  # 身体俯仰角

        # 步态参数
        self.gait_phase = 0  # 步态相位
        self.support_legs = []  # 支撑腿
        self.swing_legs = []  # 摆动腿

    async def update(self, dt: float):
        """更新控制器状态"""
        self.gait_phase = (self.gait_phase + dt / self.step_period) % 1.0
        self._update_gait()
        await self._update_leg_positions()

    def _update_gait(self):
        """更新步态"""
        # 三足步态
        if self.gait_phase < 0.5:
            self.support_legs = [0, 2, 4]  # 左前、右中、左后
            self.swing_legs = [1, 3, 5]  # 右前、左中、右后
        else:
            self.support_legs = [1, 3, 5]  # 右前、左中、右后
            self.swing_legs = [0, 2, 4]  # 左前、右中、左后

    async def _update_leg_positions(self):
        """更新腿部位置"""
        for leg_id in range(6):
            if leg_id in self.support_legs:
                await self._update_support_leg(leg_id)
            else:
                await self._update_swing_leg(leg_id)

    async def _update_support_leg(self, leg_id: int):
        """更新支撑腿位置"""
        # 计算支撑腿目标位置
        target_pos = self._calculate_support_position(leg_id)

        # 更新腿部状态
        self.leg_states[leg_id].position = target_pos
        self.leg_states[leg_id].contact = True

        # 控制舵机
        await self.move_leg(leg_id, list(target_pos), 1000)

    async def _update_swing_leg(self, leg_id: int):
        """更新摆动腿位置"""
        # 计算摆动腿轨迹
        trajectory = self._calculate_swing_trajectory(leg_id)

        # 更新腿部状态
        self.leg_states[leg_id].position = trajectory
        self.leg_states[leg_id].contact = False

        # 控制舵机
        await self.move_leg(leg_id, list(trajectory), 1000)

    def _calculate_support_position(self, leg_id: int) -> Tuple[float, float, float]:
        """计算支撑腿目标位置"""
        # 根据腿部ID和身体姿态计算目标位置
        x = self.step_length * math.cos(2 * math.pi * self.gait_phase)
        y = self.step_length * math.sin(2 * math.pi * self.gait_phase)
        z = -self.body_height

        # 应用身体姿态
        x, y, z = self._apply_body_orientation(x, y, z)

        return (x, y, z)

    def _calculate_swing_trajectory(self, leg_id: int) -> Tuple[float, float, float]:
        """计算摆动腿轨迹"""
        # 计算摆动相位
        swing_phase = (self.gait_phase - 0.5) % 1.0

        # 计算轨迹
        x = self.step_length * math.cos(2 * math.pi * swing_phase)
        y = self.step_length * math.sin(2 * math.pi * swing_phase)
        z = -self.body_height + self.step_height * math.sin(math.pi * swing_phase)

        # 应用身体姿态
        x, y, z = self._apply_body_orientation(x, y, z)

        return (x, y, z)

    def _apply_body_orientation(
        self, x: float, y: float, z: float
    ) -> Tuple[float, float, float]:
        """应用身体姿态"""
        # 创建旋转矩阵
        roll_matrix = np.array(
            [
                [1, 0, 0],
                [0, math.cos(self.body_roll), -math.sin(self.body_roll)],
                [0, math.sin(self.body_roll), math.cos(self.body_roll)],
            ]
        )

        pitch_matrix = np.array(
            [
                [math.cos(self.body_pitch), 0, math.sin(self.body_pitch)],
                [0, 1, 0],
                [-math.sin(self.body_pitch), 0, math.cos(self.body_pitch)],
            ]
        )

        # 应用旋转
        pos = np.array([x, y, z])
        pos = np.dot(roll_matrix, pos)
        pos = np.dot(pitch_matrix, pos)

        return tuple(pos)

    async def set_body_orientation(self, roll: float, pitch: float):
        """设置身体姿态"""
        self.body_roll = roll
        self.body_pitch = pitch

    async def set_motion_parameters(
        self, step_height: float, step_length: float, step_period: float
    ):
        """设置运动参数"""
        self.step_height = step_height
        self.step_length = step_length
        self.step_period = step_period

    async def emergency_stop(self):
        """紧急停止"""
        await self.motion_controller.emergency_stop()
